#ifndef REFLOW_SCEV_SCALAREVOLUTION_CANON_H
#define REFLOW_SCEV_SCALAREVOLUTION_CANON_H

#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/IR/ConstantRange.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"

namespace llvm {
class DataLayout;
class SCEVZ3Rewriter;
class Pass;
class PassRegistry;

void initializeScalarEvolutionCanonWrapperPassPass(PassRegistry &);
Pass *createScalarEvolutionCanonWrapperPassPass();

/// The main Scalar Evolution Canonicalization diver. The user is asked to use
/// this class to get canonicalized SCEV.
class ScalarEvolutionCanon
    : public SCEVVisitor<ScalarEvolutionCanon, const SCEV *> {
  Function &F;
  ScalarEvolution &SE;
  DominatorTree &DT;
  LoopInfo &LI;
  SCEVZ3Rewriter &Z3WRITER;

public:
  ScalarEvolutionCanon(Function &F, ScalarEvolution &SE, DominatorTree &DT,
                       LoopInfo &LI, SCEVZ3Rewriter &Z3WRITER)
      : F(F), SE(SE), DT(DT), LI(LI), Z3WRITER(Z3WRITER) {}

  ScalarEvolution &getSE() { return SE; }
  const ScalarEvolution &getSE() const { return SE; }

  /// Canonicalize the SCEV to a better one
  const SCEV *canonicalize(const SCEV *S) { return visit(S); }
  const SCEV *visitConstant(const SCEVConstant *S);
  const SCEV *visitTruncateExpr(const SCEVTruncateExpr *S);
  const SCEV *visitZeroExtendExpr(const SCEVZeroExtendExpr *S);
  const SCEV *visitSignExtendExpr(const SCEVSignExtendExpr *S);
  const SCEV *visitAddExpr(const SCEVAddExpr *S);
  const SCEV *visitMulExpr(const SCEVMulExpr *S);
  const SCEV *visitUDivExpr(const SCEVUDivExpr *S);
  const SCEV *visitAddRecExpr(const SCEVAddRecExpr *S);
  const SCEV *visitSMaxExpr(const SCEVSMaxExpr *S);
  const SCEV *visitUMaxExpr(const SCEVUMaxExpr *S);
  const SCEV *visitUnknown(const SCEVUnknown *S);
  const SCEV *visitCouldNotCompute(const SCEVCouldNotCompute *S);

  /// Get the canonicalized SCEV of Value V from cached ValueExprMap
  const SCEV *getExistingCanonSCEV(Value *V);

  /// Get the canonicalized SCEV of SCEV S
  const SCEV *getCanonSCEV(const SCEV *S);

  /// Get the canonicalized SCEV of Value V
  const SCEV *getCanonSCEV(Value *V);

  /// Run z3 equivalence check for each SCEV of the intructions in F
  void runZ3Check();

  /// Check if 2 SCEVs are equivalent
  void checkEquivalence(const SCEV *betterSCEV, const SCEV *S) const;

  /// Print out the non-canonicalized SCEV and the canonicalized SCEV of each
  /// SCEVable Value V Format: For V
  ///   ---> non-canonicalized SCEV
  ///   ---> canonicalized SCEV
  void print(raw_ostream &OS) const;

private:
  /// The type for ValueExprMap.
  using ValueExprMapType =
      DenseMap<Value *, const SCEV *, DenseMapInfo<Value *>>;

  /// Record the canonicalized SCEV.
  ValueExprMapType ValueExprMap;

  /// Return the DataLayout associated with the module this SCEV instance is
  /// operating on.
  const DataLayout &getDataLayout() const {
    return F.getParent()->getDataLayout();
  }

  /// Return the LLVMContext associated with the function this SCEV instance is
  /// operating on.
  LLVMContext &getContext() const {
    return F.getParent()->getContext();
  }

  /// Return the bit width of the given type
  uint64_t getTypeSizeInBits(Type *Ty) const {
    return getDataLayout().getTypeSizeInBits(Ty);
  }

  /// Sink the power of 2 constant divisor in to SCEVAddExpr or SCEV AddRecExpr.
  /// Return the operands for the UDiv part in UDs, Modulo part in Mods
  void sinkUDivInAdd(SmallVectorImpl<const SCEV *> &UDs,
                     SmallVectorImpl<const SCEV *> &Mods,
		     const SCEVNAryExpr *LHS, const SCEVConstant *RHSC);
};

class ScalarEvolutionCanonWrapperPass : public FunctionPass {
  std::unique_ptr<ScalarEvolutionCanon> SEC;

public:
  static char ID;

  ScalarEvolutionCanonWrapperPass() : FunctionPass(ID) {
    initializeScalarEvolutionCanonWrapperPassPass(
        *PassRegistry::getPassRegistry());
  }

  ScalarEvolutionCanon &getSEC() { return *SEC; }
  const ScalarEvolutionCanon &getSEC() const { return *SEC; }

  void getAnalysisUsage(AnalysisUsage &AU) const override;
  bool runOnFunction(Function &F) override;
  void releaseMemory() override;
  void print(raw_ostream &OS, const Module * = nullptr) const override;
};
} // namespace llvm

#endif // REFLOW_SCEV_SCALAREVOLUTION_CANON_H
